
import datetime
import pandas as pd
import numpy as np
from scipy import interpolate

from utils.phy import get_sensor_T, merge_data
from utils.h5data import read_h5_data, save_h5_data
import tqdm
import os

def get_st_data_2406(START_TIME, STATS,ND = 180, DT=1/48, CSVROOT='loc/P320_temp'):
    '''
    读取loc/P320_temp下的所有csv文件，并返回仪器温度数据
    '''  
    st_all = []
    t_phy_common = np.arange(0,ND,DT)
    STATS_VALID = []
    for name in tqdm.tqdm(STATS):
        file = f'{CSVROOT}/{name}.csv'
        if not os.path.exists(file):
            continue
        te_bjt, st_i = get_sensor_T(START_TIME, file)
        phy_i = merge_data(te_bjt, st_i, t_phy_common)
        st_all.append(phy_i)
        STATS_VALID.append(name)
    
    return STATS_VALID,t_phy_common,st_all


if __name__ == '__main__':

    from myglobal import get_info
    from utils.loc import get_distance, sort_data_by_distance
    
    DATE='254C'
    LINE='254C'

    marker = f'{DATE}.{LINE}'
    CASVROOT = 'loc/254C_temp'
    D_info = get_info(DATE)
    START_TIME = D_info['START_TIME']
    LINES = D_info['LINES']
    s_info = D_info['s_info']
    STATS = LINES[LINE]
    ND = D_info['NDAYS']
    
    STATS,te_bjt,phy_all = get_st_data_2406(START_TIME,STATS, ND, CSVROOT=CASVROOT)

    print([phy.std() for phy in phy_all])

    x_used = []
    STATS_used = []
    phy_used = []
    for i,name in enumerate(STATS):
        xi, _,_ = get_distance(s_info=s_info,name1='P264',name2=name, S2N=False)
        if phy_all[i].std()<=100:
            x_used.append(xi)
            STATS_used.append(name)
            phy_used.append(phy_all[i])
        
    x,STATS = sort_data_by_distance(x_used,NAMES=STATS_used)
    phy_all = np.array(phy_used)

    x = np.array(x)
    DX = np.diff(x).mean()
    nx = len(x)

    phy_all= np.array(phy_all)
    nt = len(te_bjt)
    phy_mean = phy_all.mean(axis=0)
    slopes = []
    R2 = []
    for i in range(nt):
        phy_i = phy_all[:,i]
        coffs = np.polyfit(x, phy_i, 1)
        k,b = coffs
        # 计算拟合程度
        R2.append(1-np.sum((phy_i-np.polyval(coffs,x))**2)/np.sum((phy_i-phy_i.mean())**2))
        slopes.append(-k*1000)
    slopes = np.array(slopes)
    R2 = np.array(R2)
    save_h5_data(f'./loc/{marker}.sensorT.h5',
                 data_in_dict={'te_bjt':te_bjt,
                               'slopes':slopes,
                               'R2':R2,
                               'phy_all':phy_all,
                               'phy_mean':phy_mean,
                               'STATS':np.array(STATS, dtype='S'),
                               'START_TIME':START_TIME},
                 attrs_dict={'marker':marker,'START_TIME':START_TIME})

    from pylab import figure, rcParams, plt
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '8'
    
    fig = figure(figsize=(8, 8), dpi=500)
    ax1 = plt.subplot(311)
    SCALE= DX/phy_all.std()
    for i in range(nx):
        phy_i = phy_all[i,:]
        phy_i = (phy_i-phy_i.mean())*SCALE
        line, = plt.plot(te_bjt, phy_i+x[i], lw=0.5, color='k')

    plt.xlim([10,20])
    plt.xlabel('day (BJT)')
    plt.ylabel('x(m)',color='k')

    ax2 = plt.subplot(312)

    line, = plt.plot(te_bjt, phy_mean, lw=0.5, color='k')
    plt.xlim([10,20])
    plt.xlabel('day (BJT)')
    plt.ylabel('T_mean',color='tab:red')

    ax3 = plt.subplot(313)

    line, = plt.plot(te_bjt, slopes, lw=0.5, color='k')
    plt.ylabel('dT/km',color='tab:red')
    ax4 = plt.twinx(ax3)
    line, = plt.plot(te_bjt, R2, lw=0.5, color='tab:green')
    plt.ylabel('R2',color='tab:green')

    plt.xlim([10,20])
    plt.xlabel('day (BJT)')

    plt.savefig(f'./loc/{marker}.sensorT.png')


